package com.ting.test.algorithms.图;

import java.util.*;

/**
 * 1）总是从权值最小的边开始考虑，依次考察权值依次变大的边
 * 2）当前的边要么进入最小生成树的集合，要么丢弃
 * 3）如果当前的边进入最小生成树的集合中不会形成环，就要当前边
 * 4）如果当前的边进入最小生成树的集合中会形成环，就不要当前边
 * 5）考察完所有边之后，最小生成树的集合也得到了
 * <p>
 * 具体的解题思路如下：
 * <p>
 * 每次添加的是一条边
 * <p>
 * 1. 将节点按照长度（权重）从小到大进行排序
 * 2. 利用并查集的思路，把每一个节点作为一个集合，如果一条边对应的节点A和B已经在同一个集合中，则再添加该边的话就会形成环，因此不能要该边，否则不会，可以要当前边。
 * 3. 当所有边都考察完时，最小生成树就出来了。
 */
public class 最小生成树之K算法 {

    public static void main(String[] args) {
        List<Edge> result = minTreeByK(new Graph());
    }

    private static List<Edge> minTreeByK(Graph graph) {
        List<Edge> result = new ArrayList<>();
        UnionSet unionSet = new UnionSet();
        unionSet.makeUnionSet(graph);
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<Edge>(new MyComparator());
        //吧所有的边从小到大排序
        for (Edge edge : graph.edges) {
            priorityQueue.add(edge);
        }
        while (!priorityQueue.isEmpty()) {
            Edge poll = priorityQueue.poll();
            //判断是否属于同一个集合，如果是则表明该点已经添加过了
            if (unionSet.isSameSet(poll.getFrom(), poll.getTo())) {
                continue;
            }
            //如果不是，则合并集合
            unionSet.merge(poll.getFrom(), poll.getTo());
            result.add(poll);
        }
        return result;
    }

    static class UnionSet {
        HashMap<Node, Node> parentMap = new HashMap<>();
        HashMap<Node, Integer> sizeMap = new HashMap<>();

        public void makeUnionSet(Graph graph) {
            if (graph == null) {
                return;
            }

            for (Node node : graph.nodes.values()) {
                parentMap.put(node, node);//初始化集合
                sizeMap.put(node, 1);//初始化集合
            }
        }

        public boolean isSameSet(Node nodeA, Node nodeB) {
            return findParent(nodeA) == findParent(nodeB);

        }

        public void merge(Node nodeA, Node nodeB) {
            Node parentA = findParent(nodeA);
            Node parentB = findParent(nodeB);

            if (parentA.equals(parentA)) {
                return;
            }

            int sizeA = sizeMap.get(parentA);
            int sizeB = sizeMap.get(parentB);
            if (sizeA > sizeB) {
                parentMap.put(parentB, parentA);
                sizeMap.remove(parentB);
                sizeMap.put(parentA, sizeA + sizeB);
            } else {
                parentMap.put(parentA, parentB);
                sizeMap.remove(parentA);
                sizeMap.put(parentB, sizeA + sizeB);
            }
        }

        private Node findParent(Node nodeB) {
            Stack<Node> stack = new Stack<>();
            Node parent = nodeB;
            while (!parent.equals(parentMap.get(parent))) {
                stack.push(parent);//入栈
                parent = parentMap.get(parent);

            }
            //将沿途经过的每个点的parent都设置为最顶层的parent,加快查询速度
            while (!stack.isEmpty()) {
                Node pop = stack.pop();
                parentMap.put(pop, parent);
            }

            return parent;
        }

    }

    static class MyComparator implements Comparator<Edge> {
        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.getWeight() - o2.getWeight();
        }
    }
}
